from models.resnet import *
import torch.nn as nn
import torch
from models.globalNet import globalNet
from models.refineNet import refineNet
# from models.simple_pose import get_pose_net
from models.Unet import UNet
from models.Unet_resnet import unet_resnet_model
from models.D_LinkNet import model_DinkNet34
from models.hourglass import PoseNet
from models.single_model import single_unet_resnet_model
from models.Unet_resnet_GN import unet_resnet_model_GN
from models.Hrnet.pose_hrnet import get_hr_net
from models.hrseg.seg_hrnet import get_hrseg_model
from models.RSN18.network import model_RSN18
from models.efficientunet import get_efficientunet_b0,get_efficientunet_b3,get_efficientunet_b4
__all__ = ['CPN50', 'CPN101']

class CPN(nn.Module):
    def __init__(self, resnet, output_shape, num_class, pretrained=True):
        super(CPN, self).__init__()
        channel_settings = [2048, 1024, 512, 256]
        self.resnet = resnet
        self.global_net = globalNet(channel_settings, output_shape, num_class)
        self.refine_net = refineNet(channel_settings[-1], output_shape, num_class)

    def forward(self, x):
        res_out = self.resnet(x)
        global_fms, global_outs = self.global_net(res_out)
        refine_out = self.refine_net(global_fms)

        return global_outs, refine_out

def CPN50(out_size,num_class,pretrained=True):
    res50 = resnet50(pretrained=pretrained)
    model = CPN(res50, output_shape=out_size,num_class=num_class, pretrained=pretrained)
    return model

def CPN101(out_size,num_class,pretrained=True):
    res101 = resnet101(pretrained=pretrained)
    model = CPN(res101, output_shape=out_size,num_class=num_class, pretrained=pretrained)
    return model

def model_cpn50(image_size = (320,512)):
    model = CPN50(out_size=image_size, num_class=7, pretrained=False).cuda()
    model_state_dict = model.state_dict()
    pre_trained_state_dict = torch.load('../pre_trained_models/CPN50_384x288.pth')['state_dict']
    pre_trained_state_dict = {k.replace('module.',''):v for k,v in pre_trained_state_dict.items() if 'predict' not in k }
    model_state_dict.update(pre_trained_state_dict)
    model.load_state_dict(model_state_dict)
    return model

def model_simple_pose_res101(pretrain = False,num_class = 7):
    model = get_pose_net(num_layers=101,num_class = num_class).cuda()
    if pretrain:
        model_state_dict = model.state_dict()
        pre_trained_dict = torch.load('../pre_trained_models/pose_resnet_101_384x288.pth')
        pre_trained_dict = {k: v for k, v in pre_trained_dict.items() if 'final' not in k}
        model_state_dict.update(pre_trained_dict)
        model.load_state_dict(model_state_dict)
    return model

def model_simple_pose_res18(pretrain=False):
    model = get_pose_net(num_layers=18).cuda()
    if pretrain:
        model_state_dict = model.state_dict()
        pre_trianed_dict = torch.load('../pre_trained_models/resnet18.pth')
        pre_trianed_dict = {k: v for k, v in pre_trianed_dict.items() if 'fc' not in k}
        model_state_dict.update(pre_trianed_dict)
        model.load_state_dict(model_state_dict)
    return model

def model_unet():
    model = UNet(n_class=7).cuda()
    return model

def model_resnet_unet(layer,num_classes,use_aspp=False,pre_train=False):
    model = unet_resnet_model(layers=layer,num_classes=num_classes,use_aspp=use_aspp,pre_train=pre_train)
    return model

def model_resnet_unet_GN(layer,num_classes,use_aspp=False,pre_train=False):
    model = unet_resnet_model_GN(layers=layer,num_classes=num_classes,use_aspp=use_aspp,pre_train=pre_train)
    return model

def model_DLinknet34(pre_trained = True,classify = False):
    model = model_DinkNet34(pre_trained = pre_trained,classify = classify)
    return model

def model_posenet(pre_trained = True):
    model = PoseNet().cuda()
    if pre_trained:
        model_state_dict = model.state_dict()
        pretrain_state_dict = torch.load('../../pre_trained_models/hourglass_2.pt')['state_dict']
        pretrain_state_dict = {k.replace('model.module.', ''): v for k, v in pretrain_state_dict.items() if
                               ('out' not in k and 'pred' not in k)}
        model_state_dict.update(pretrain_state_dict)
        model.load_state_dict(model_state_dict)
    return model

def model_single_unet_res34(pre_trained = False):
    model = single_unet_resnet_model(layers='resnet34',pre_train=pre_trained)
    return model

def model_hrnet(pretrained,num_class,upconv = False):
    model = get_hr_net(pretrained=pretrained,num_classes=num_class,upconv = upconv).cuda()
    return model
    
def model_small_seghrnet():
    model = get_hrseg_model(cfg_dir='./models/hrseg/cityscapes/seg_hrnet_w18_small_v2_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml')
    model.cuda()
    return model

def model_rsn18(num_class):
    model = model_RSN18(num_class=num_class)
    model = model.cuda()
    return model

def model_Unet_efficient_b0(num_class):
    model = get_efficientunet_b0(out_channels=num_class,pretrained=False)
    model = model.cuda()
    return model
def model_Unet_efficient_b3(num_class):
    model = get_efficientunet_b3(out_channels=num_class,pretrained=False)
    model = model.cuda()
    return model
def model_Unet_efficient_b4(num_class):
    model = get_efficientunet_b4(out_channels=num_class,pretrained=False)
    model = model.cuda()
    return model
if __name__ == '__main__':

    model = model_resnet_unet_GN('resnet34',num_classes=2)
    print(model)